
#Author: Mike Schladt - 2014
#Program loads malware samples from inbox into cuckoo

import os
import sys
import signal
import json
import re
import threading 
import time
import lib.mipreport as report
from lib.mipsample import MipSample

def signal_handler(signal, frame):
    sys.exit(0) #Process will terminate after all threads complete

signal.signal(signal.SIGINT, signal_handler)

defaultInbox = os.environ['HOME']+"/samples/inbox/"

try: 
    inbox = sys.argv[1]
except:
    inbox = defaultInbox

print ("Master Controller Starting...")
print ("Searching for malware files in: " + inbox)

#find all files in inbox (recurses directories)
for root, dirs, files in os.walk(inbox):
    for file in files:
        try:
            #initalize sample -- find file attributes
            sample_path = os.path.join(inbox,file)
            sample = MipSample(sample_path)
            
            #store sample attribute to database
            try :        
                sample.store_db()

            except Exception as e:
                print e

            #sumbit sample to cuckoo REST service
            try : 
                sample.submit()    

            except Exception as e :
                print e
        
            print "loading: " + sample.file + " SID: " + str(sample.sid) + \
                " Task: " + str(sample.task_id)  

            #process cuckoo report as new thread 
            t = threading.Thread(target=report.process_cuckoo_report, args=(sample,))
            t.start()

            #grab the status of cuckoo
            status = report.get_cuckoo_status()
            
            print "Cuckoo Pending: " + str(status['tasks']['pending'])
            print "Cuckoo Running: " + str(status['tasks']['running'])
            print "Cuckoo Completed: " + str(status['tasks']['completed'])
            print "Cuckoo Reported: " + str(status['tasks']['reported'])
            print "Cuckoo Total: " + str(status['tasks']['total'])
            print "Machines available: " + str(status['machines']['available']) + \
                " of " + str(status['machines']['total'])
          
            #wait for free machines
            while status['machines']['available'] < 1 : 
                print "Waiting for free machine"
                print "Machines available: " + str(status['machines']['available'])
                status = report.get_cuckoo_status()                
                time.sleep(6)

            #wait if reporting is bogged down
            while status['tasks']['completed'] > 16 :
                print "Waiting for reporting to catch up"
                print "Waiting on " + str(status['tasks']['completed']) + " reports"
                status = report.get_cuckoo_status()
                time.sleep(6) 


            time.sleep(6) #give machine time to load before next thread
        except Exception as e:
            print e

#wait for threads to finish
while threading.active_count() > 1:
    pass





